# -*- coding:utf-8 -*-
try:
    from Crypto.Cipher import AES
    from Crypto import Random
except ImportError:
    raise ImportError("Can't find package:Please install package 'pycryptodome'.")

class AESCrypt(object):
    '''一个用于AES加密的类。目前仅支持ECB方法加解密
    '''
    def __init__(self,password = None):
        '''构造一个加密器
            密码留空将生成一个随机密码
        '''
        if password is None:
            password = Random.get_random_bytes(32)
        else:
            password = str(password)
        self.__password = self._complete(password)

    def _complete(self,text,encoding = 'utf-8'):
        '''用于把不足16位的文本编码后补足'''
        if isinstance(text,str):
            text = text.encode(encoding = encoding)
        while len(text) % 16 != 0:
            text += b'\x00'           #补位
        return text

    def encrypt(self,text,encoding = 'utf-8'):
        '''进行加密。text必须是一个能转换为字符串的类型或byte'''
        if isinstance(text,bytes):
            text = self._complete(text)
        else:
            text = self._complete(str(text),encoding = encoding)
        aes = AES.new(self.__password,AES.MODE_ECB)
        return aes.encrypt(text)

    def decrypt(self,text,mode = 'str',encoding = 'utf-8'):
        '''进行解密。text必须是AES加密的密文
            mode可设为'str'或'byte',设为'str'将自动解码后返回
        '''
        aes = AES.new(self.__password,AES.MODE_ECB)
        text = aes.decrypt(text).rstrip(b'\x00')
        if mode == 'str':
            return text.decode(encoding = encoding)
        elif mode == 'byte':
            return text
        else:
            raise ValueError('mode必须是"str"或"byte"中的一种！')

